from utils.log_helper import log_init
from config.config import cfg
from utils.data_helper import LoadCoupletDataset
from loguru import logger

if __name__ == "__main__":
    log_init(filename="LoadCoupletDataset", level="DEBUG")
    data_loader = LoadCoupletDataset(cfg.train_corpus_file_paths,
                                     batch_size=cfg.batch_size,
                                     top_k=cfg.top_k)
    train_iter, test_iter = data_loader.load_train_val_test_data(cfg.train_corpus_file_paths,
                                                                 cfg.test_corpus_file_paths)
    logger.debug(f'PAD_IDX = {data_loader.PAD_IDX}')
    for src, tgt in train_iter:
        tgt_input = tgt[:-1, :]
        tgt_out = tgt[1:, :]
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = data_loader.create_mask(src, tgt_input)
        logger.debug(f'src = {src}')
        logger.debug(f'tgt = {tgt}')

        # [in_tensor_len, batch_size]
        logger.debug(f'src shape: {src.shape}')
        # [in_tensor_len, batch_size]
        logger.debug(f'tgt shape: {tgt.shape}')
        logger.debug(f'src input shape: {src.shape}')
        logger.debug(f'src_padding_mask shape (batch_size, src_len): {src_padding_mask.shape}')
        logger.debug(f'tgt input shape: {tgt_input.shape}')
        logger.debug(f'tgt_padding_mask shape: (batch_size, tgt_len) {tgt_padding_mask.shape}')
        logger.debug(f'tgt output shape: {tgt_out.shape}')
        logger.debug(f'tgt_mask shape (tgt_len,tgt_len):  {tgt_mask.shape}')
        break
